import pandas as pd
import requests
import time
from tqdm import tqdm
import json

# --- 1. 参数配置 ---

#在这里放入APIkey
AMAP_API_KEY = "865cb93a62736536b5b1146de0328d0d"

# 输入和输出文件名
INPUT_CSV_PATH = "wuhan_price_part_2.csv"  # ！！！！！！注意改成自己处理的部份
OUTPUT_CSV_PATH = "wuhan_price_data_with_coords_2_final.csv"  # !!!!!!注意改成自己处理的部份

# API请求的固定参数
CITY_NAME = "武汉"  # 明确指定在“武汉市”内搜索，确保结果的准确性
API_URL = "https://restapi.amap.com/v3/geocode/geo"  # 高德地理编码API

# 速率控制：每次API请求之间的延迟时间（秒）
REQUEST_DELAY = 0.1


# --- 2. 定义核心功能函数 ---

def get_coordinates(address: str, api_key: str, city: str):
    """
    使用高德地图API获取单个地址的经纬度坐标。

    Args:
        address (str): 需要查询的小区名称或详细地址。
        api_key (str): 您的高德Web服务API Key。
        city (str): 查询的目标城市，用于限定搜索范围。

    Returns:
        tuple: (经度, 纬度) or (None, None) 如果查询失败。
    """
    params = {
        'key': api_key,
        'address': address,
        'city': city
    }
    try:
        response = requests.get(API_URL, params=params)
        response.raise_for_status()  # 如果请求失败 (例如 404, 500), 抛出异常

        result = response.json()

        if result['status'] == '1' and result['geocodes']:
            # 提取第一个匹配结果的坐标
            location_str = result['geocodes'][0]['location']
            longitude, latitude = location_str.split(',')
            return float(longitude), float(latitude)
        else:
            # API返回成功但没有找到地理编码
            # print(f"警告: 地址 '{address}' 在城市 '{city}' 未找到坐标。API返回信息: {result.get('info')}")
            return None, None

    except requests.exceptions.RequestException as e:
        print(f"错误: 请求API时发生网络错误: {e}")
        return None, None
    except (KeyError, IndexError, json.JSONDecodeError) as e:
        print(f"错误: 解析API响应时出错: {e}")
        return None, None


# --- 3. 主程序逻辑 ---

def main():
    """
    主程序，用于读取CSV，处理数据并保存结果。
    """
    print("程序开始...")

    # 尝试读取CSV文件
    try:
        df = pd.read_csv(INPUT_CSV_PATH, encoding='utf-8-sig')
    except FileNotFoundError:
        print(f"错误: 输入文件 '{INPUT_CSV_PATH}' 未找到。请检查文件名和路径。")
        return
    except Exception as e:
        print(f"错误: 读取CSV文件时发生错误: {e}")
        return

    # 检查DataFrame是否为空
    if df.empty:
        print("警告: 读取的CSV文件为空，程序已退出。")
        return

    # ==================================================================
    # --- 新增：试运行模块 ---
    print("\n--- 开始单次试运行 ---")

    # 找到CSV文件中第一个非空的小区名称用于测试
    first_community_name = df['小区'].dropna().iloc[0]

    if first_community_name:
        print(f"正在测试地址: '{first_community_name}'")
        # 调用API函数获取坐标
        test_lon, test_lat = get_coordinates(first_community_name, AMAP_API_KEY, CITY_NAME)

        if test_lon is not None and test_lat is not None:
            print(f"✅ 测试成功! 获取到坐标: 经度={test_lon}, 纬度={test_lat}")
        else:
            print(f"❌ 测试失败。未能获取到坐标，请检查API Key、网络或地址名称。")
    else:
        print("警告: 文件中未找到有效的小区名称进行测试。")

    print("--- 试运行结束 ---\n")
    # ==================================================================

    # 初始化新的经纬度列
    df['longitude'] = None
    df['latitude'] = None

    print(f"成功读取 {len(df)} 条数据。即将开始批量地理编码处理...")

    # 使用tqdm创建进度条，并迭代处理每一行数据
    # total=df.shape[0] 设置进度条的总长度
    for index, row in tqdm(df.iterrows(), total=df.shape[0], desc="小区坐标获取中"):
        community_name = row['小区']

        if pd.notna(community_name):
            # 调用函数获取坐标
            lon, lat = get_coordinates(community_name, AMAP_API_KEY, CITY_NAME)

            # 将获取到的坐标填充到DataFrame中
            if lon is not None and lat is not None:
                df.loc[index, 'longitude'] = lon
                df.loc[index, 'latitude'] = lat

        # 停顿，以满足API速率限制
        time.sleep(REQUEST_DELAY)

    # 保存处理完成的数据到新的CSV文件
    # index=False 表示不将DataFrame的索引写入文件
    # encoding='utf-8-sig' 确保中文在Excel中能正确显示
    df.to_csv(OUTPUT_CSV_PATH, index=False, encoding='utf-8-sig')

    print("\n处理完成！")
    print(f"带有坐标的数据已保存到: {OUTPUT_CSV_PATH}")

    # 统计成功和失败的数量
    success_count = df['longitude'].notna().sum()
    total_count = len(df)
    print(f"总计 {total_count} 条数据，成功获取了 {success_count} 个坐标，失败 {total_count - success_count} 个。")


if __name__ == '__main__':
    main()